rm(list = ls())
library(tidyverse)
library(estimatr)
library(grid)
library(ri2)
library(reshape2)
library(bandit)
library(stargazer)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
set.seed(10027)
source('fn.R')

load('data.rdata')

mod  <- lm_robust(Y ~ arm - 1, data = data, weights = ipw)
true <- coef(mod)
iter <- 10000

if (file.exists('Output/TF2_simulation.rdata')){
  load('Output/TF2_simulation.rdata')
}else {
  simlist <- list()
  for (p in 2:5){
    simlist[[paste0('p', p)]] <- simulate(periods = p, n = round(2000/p), probs = true, iter = iter, static = FALSE, ppmat = TRUE)
  }
}

sumlist <- list()
for (p in 2:5){
  sumlist[[paste0('p', p)]] <- result_sum(data = as.data.frame(simlist[[paste0('p', p)]][['d_fit']])) %>% select(-est, -bias)
}

result <- sumlist$p2 %>%
  left_join(sumlist$p3 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.3'))) %>%
  left_join(sumlist$p4 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.4'))) %>%
  left_join(sumlist$p5 %>% rename_at(vars(matches('best|rmse|coverage')), ~paste0(.x, '.5'))) %>%
  select(-true)

stargazer::stargazer(result, summary = F, rownames = F)


